#! /usr/bin/python
# -*- coding: UTF-8 -*-
# 添加OC垃圾代码
import os,sys
import random
import string
import re
import md5
import time
import json
import shutil
import hashlib 
import time
import argparse
import codecs
from biplist import *

import sys 
reload(sys) 
sys.setdefaultencoding("utf-8")

script_path = None
target_ios_folder= None
total_class_num = 60
match_rule = None
Class_name_list= None
Class_name_array =None
insertAPI_name_list = None
word_name_list = None
#确保添加的函数不重名
funcname_set = set()
g_funcname_set = set()
headerCollectFileName = None
totalSize = 0
fileNum = 0
dirNum= 0

def initAddNativeCode():
    global funcname_set,g_funcname_set
    funcname_set = set()
    g_funcname_set = set()

    global script_path
#    script_path = os.path.abspath(os.path.dirname(__file__))
    script_path = "/Users/k/Desktop/7YAO/Mac_Obfunscation/Mac_Obfunscation/addNativeScriptForOC/"

    #匹配规则，路径包含path_include且不包含path_exclude的才会创建对应类型的文件
    global match_rule
    match_rule = {
        ".h": {
    #        "path_include": os.path.sep + "res",
            "path_include": os.path.sep + "Bundle.bundle"
        },
        ".h": {
            # "path_include": "src",
    #        "path_exclude": os.path.sep + "res",
            "path_exclude": os.path.sep + "Bundle.bundle"
        },
    }

    #方法列表，用以随机方法
    modify_Class_name()
    global Class_name_list
    with open(os.path.join(script_path, "./Class_list.json"), "r") as fileObj:
        Class_name_list = json.load(fileObj)

    global Class_name_array
    with open(os.path.join(script_path, "./Class.json"), "r") as fileObj:
        Class_name_array = json.load(fileObj)
        # print  ('dasfaf')

    #单词列表，用以随机名称
    global insertAPI_name_list
    with open(os.path.join(script_path, "./insertAPIName.plist"), "r") as fileObj:
        insertAPI_name_list = readPlist(fileObj)

    #单词列表，用以随机名称
    global word_name_list
    with open(os.path.join(script_path, "./word_list.json"), "r") as fileObj:
        word_name_list = json.load(fileObj)
    # print('addNativeCode script_path:',script_path)
    pass

#获取随机字符串
def generate_array_className(number):
    global Class_name_list
    return Class_name_list[number]


def generate_className_inArray(number):
    global Class_name_array
    return Class_name_array[number]

#获取随机字符串
def generate_random_ClassName():
    global Class_name_list
    return random.choice(Class_name_list)

#获取一个类名
def getClassName(number):
    return generate_array_className(number) + 'ViewController'

#oc代码头文件函数声明
def getOCHeaderFuncText(number):
    funcName = 'pushTo' + generate_array_className(number) + 'VC'
    text = "\n- (void)%s" %(funcName)
    return text


def getOCHeaderContent(number):
    funcName = 'pushTo' + generate_className_inArray(number) + 'VC'
    text = "\n- (void)%s" %(funcName)
    return text

#oc代码函数实现模板
def getOCFuncText(header_text):
    global insertAPI_name_list
    arg1 = random.choice(insertAPI_name_list)
    text = [
        header_text + "\n",
        "{\n",
        "\t %s  \n" %(arg1),
        "}\n"
    ]
    return "".join(text)


def getOCFunctionContent(header_text):
    global insertAPI_name_list
    arg1 = random.choice(insertAPI_name_list)
    text = [
        header_text + "\n",
        "{\n",
        "\t %s  \n" %(arg1),
        "}\n"
    ]
    return "".join(text)


#oc代码以@end结尾，在其前面添加text
def appendTextToOCFile(file_path, text):
    with open(file_path, "r") as fileObj:
        old_text = fileObj.read()
        fileObj.close()
        end_mark_index = old_text.rfind("@end")
        if end_mark_index == -1:
            print "\t非法的结尾格式: " + file_path
            return
        new_text = old_text[:end_mark_index]
        new_text = new_text + text + "\n"
        new_text = new_text + old_text[end_mark_index:]

    with open(file_path, "w") as fileObj:
        fileObj.write(new_text)

#处理单个OC文件，添加垃圾函数。确保其对应头文件存在于相同目录
def addDealFunctionWithOCFile(filename, file_path):

    end_index = file_path.rfind(".")
    pre_name = file_path[:end_index]
    header_path = pre_name + ".h"
    if not os.path.exists(header_path):
        print "\t相应头文件不存在：" + file_path
        return

    global Class_name_array
    new_func_num =  len(Class_name_array)
    for i in range(new_func_num):
        header_text = getOCHeaderContent(i)
        result = filename.split('ViewController.m')[0] in header_text
        if result == True:
            print  '此时包含'
        else:
            appendTextToOCFile(header_path, header_text + ";\n")
            funcText = getOCFunctionContent(header_text)
            appendTextToOCFile(file_path, funcText)

#扫描parent_folder，添加垃圾函数，处理了忽略列表
def addOCFunctionsToTrashFile(parent_folder):
    for parent, folders, files in os.walk(parent_folder):
        for file in files:
            if file.endswith(".m") or file.endswith(".mm"):
                addDealFunctionWithOCFile(file, os.path.join(parent, file))


#新创建的垃圾文件header模板
def getOCHeaderFileText(class_name):

    propertyStr = ''

    for header_text in Class_name_list:
        
        result = class_name.split('ViewController')[0] in header_text

        if result == True:
            print  '此时包含'
        else:
          propertyStr += "@property (strong,nonatomic)NSString * %s; \n" %(header_text)

    text = [
        "//  UnityAppController.h \n",
        "//  KXNewVersionCheck \n",
        "//  Created by Apple. \n",
        "//  Copyright © 2018年 All rights reserved. \n\n",
        "#import <Foundation/Foundation.h> \n",
        "#import <UIKit/UIKit.h> \n\n",
        "#import \"UnityAppController.h\" \n\n",
        "@interface %s : UnityAppController \n" %(class_name),
        "%s" %(propertyStr),
        "\n\n@end"
    ]
    return string.join(text)


# 新创建的垃圾文件header模板
def getOCHeaderFileContent(class_name):
    propertyStr = ''

    for header_text in Class_name_array:

        result = class_name.split('ViewController')[0] in header_text

        if result == True:
            print  '此时包含'
        else:
            propertyStr += "@property (strong,nonatomic)NSString * %s; \n" % (header_text)

    text = [
        "//  UnityAppController.h \n",
        "//  KXNewVersionCheck \n",
        "//  Created by Apple. \n",
        "//  Copyright © 2018年 All rights reserved. \n\n",
        "#import <Foundation/Foundation.h> \n",
        "#import <UIKit/UIKit.h> \n\n",
        "#import \"UnityAppController.h\" \n\n",
        "@interface %s : UnityAppController \n" % (class_name),
        "%s" % (propertyStr),
        "\n\n@end"
    ]
    return string.join(text)




#新创建的垃圾文件mm模板
def getOCMMFileText(class_name):
    text = [
        "//  UnityAppController.m\n",
        "//  KXNewVersionCheck\n",
        "//  Created by Apple.\n",
        "//  Copyright © 2018年 All rights reserved.\n\n",
        '#import "%s.h"\n\n' %(class_name),
        '@interface %s()\n' %(class_name),
        "\n@end\n\n"
        "@implementation %s\n" %(class_name),
        "\n\n@end"
    ]
    return string.join(text)

def addOCFileContent(parent_folder):
    file_list = []
    target_folder = os.path.join(parent_folder,  getOneName())
    if os.path.exists(target_folder):
        shutil.rmtree(target_folder)
    os.mkdir(target_folder)

    # global Class_name_array
    # file_num = len(Class_name_array)
    file_num = random.randint(2,6)

    end = "/"
    string2 = target_folder[target_folder.rfind(end):]
    if  len(string2) > 0:
        string2 = string2[1:]
    for i in range(file_num):
        global g_funcname_set
        if len(g_funcname_set)>= total_class_num :
            break
        file_name =   string2  +  getOneName() + "ViewController"
        g_funcname_set.add("#import \"" + file_name + ".h\"")
        header_text = getOCHeaderFileContent(file_name)
        full_path = os.path.join(target_folder, file_name + ".h")

        # global target_ios_folder
        # full_path = os.path.join(target_ios_folder, "target_ios")
        # if not os.path.exists(full_path):
        #     os.mkdir(full_path)
        # full_path = os.path.join(full_path, file_name + ".h")

        with open(full_path, "w") as fileObj:
            fileObj.write(header_text)
            fileObj.close()

        mm_text = getOCMMFileText(file_name)
        full_path = os.path.join(target_folder, file_name + ".m")
        with open(full_path, "w") as fileObj:
            fileObj.write(mm_text)

#生成一个指定长度的随机字符串
def generate_random_str():
    random_str = ''
    base_str = 'ABCDEFGHIGKLMNOPQRSTUVWXYZabcdefghigklmnopqrstuvwxyz'
    length = len(base_str) - 1
    for i in range(8):
        random_str += base_str[random.randint(0, length)]
    return random_str



#清空类名、重新生成新的
def modify_Class_name():
    
    global script_path
    with open(os.path.join(script_path,"./Class_list.json"), "r+") as f:
        read_data = f.read()
        f.seek(0)
        f.truncate()   #清空文件

    global total_class_num

    class_name_arr = []

    for index in range(total_class_num):

        random_str_class = generate_random_str()
        class_name_arr.append(random_str_class)

#    global script_path
    with open(os.path.join(script_path,"./Class_list.json"), "w") as f:
        json.dump(class_name_arr,f)
        f.closed


#添加单个文件
def addSingleFile(file_path):
    global target_ios_folder
    global Class_name_list
    print "add file " + file_path.replace(target_ios_folder, "")
    _, file_type = os.path.splitext(file_path)
    if file_path in "/sourceFile":
        return
    if file_type == ".h":
        with open(file_path, "w") as fileObj:
            func_num = random.randint(2, 5)
            for j in range(0, func_num):
                func_text = getOCFucText()
                fileObj.write(func_text)




            fileObj.close()


#获取一个随机名称
def getOneName():
    global word_name_list
    return random.choice(word_name_list)

# 获取OC垃圾方法
def getOCFucText():
    global funcname_set
    new_func_name = getOneName()
    while new_func_name in funcname_set:
        new_func_name = getOneName()
    funcname_set.add(new_func_name)

    argv_name = getOneName()
    text = [
        '\n- (void)' + new_func_name + '{\n',
        '\n',
        '\n',
        '}\n'
    ]
    return string.join(text)

def addOCTrashFile(parent_folder, level, min_file_num = 0):

    global g_funcname_set
    if len(g_funcname_set)>= total_class_num :
        return
        pass
    global match_rule
    create_folder_list = []

    for parent, folders, files in os.walk(parent_folder):
        #创建文件数量
        new_file_num = random.randint(len(files) / 2, len(files)) + min_file_num
        for i in range(0, new_file_num):
            # file_path = os.path.join(parent, getOneName() + target_file_type)
            addOCFileContent(parent)

        #防止创建太多层的文件夹
        if level > 1:
            continue
        #创建文件夹数量
        # new_fold_num = random.randint(len(files) /  2, len(files))
        new_fold_num = random.randint(1, 44)
        for i in range(0, new_fold_num):
            target_folder = os.path.join(parent, getOneName())
            #为了不阻断os.walk,延后创建文件夹
            create_folder_list.append(target_folder)

    for folder_path in create_folder_list:
        try:
            # print "create folder " + folder_path.replace(target_folder, "")
            os.mkdir(folder_path)
            addOCTrashFile(folder_path, level + 1, random.randint(1, 3))
        except Exception as e:
            print e

def DeleteEmptyDir(path):
    global totalSize
    global fileNum
    global dirNum
    for lists in os.listdir(path):
        sub_path = os.path.join(path, lists)
        if os.path.isfile(sub_path):
            fileNum = fileNum+1                      
            totalSize = totalSize+os.path.getsize(sub_path)  
        elif os.path.isdir(sub_path):
            if not os.listdir(sub_path):
                os.rmdir(sub_path)
            else:
                dirNum = dirNum+1                      
                DeleteEmptyDir(sub_path) 

def main():

    global target_ios_folder
    print "开始创建oc文件到trash目录",target_ios_folder
    addOCTrashFile(target_ios_folder, 0,1)
    # print "\n开始添加oc方法"
    addOCFunctionsToTrashFile(target_ios_folder)

    global totalSize
    global fileNum
    global dirNum
    totalSize = 0
    fileNum = 0
    dirNum= 0
    DeleteEmptyDir(target_ios_folder)
    print('dirNum fileNum,totalSize,g_funcname_setLen:',dirNum,fileNum,totalSize/float(1048576),len(g_funcname_set))
    
    all_header_text = "\n".join(g_funcname_set)
    global headerCollectFileName
    headerFileName = "{0}.h".format(headerCollectFileName)
    with open(os.path.join(target_ios_folder, headerFileName), "a+") as fileObj:
        fileObj.write("\n"+all_header_text)
        fileObj.close()
    print "\n addNativeCodeForOC Finished"


def StartAddNativeCodeForOC(allCount,trashTargetDir,collectFileName):
    #生成垃圾文件个数
    global total_class_num
    total_class_num = allCount

    #头文件汇总名字
    global headerCollectFileName
    headerCollectFileName = collectFileName

    #垃圾代码临时存放目录
    global target_ios_folder
    target_ios_folder = trashTargetDir
    main()
    pass



